import cv2
import numpy as np


def add_seg_to_img(img, seg, class_color_dict):
    # add_mask_to_img(image, seg, {0:[0,0,0], 1:[255,0,0],2:[0,255,0]})
    ret_img = img.copy()
    for cl, color in class_color_dict.items():
        seg_single_class = seg.copy()
        tmp_image = np.zeros(img.shape, dtype=np.uint8)
        tmp_image[:, :, 0] = color[0]
        tmp_image[:, :, 1] = color[1]
        tmp_image[:, :, 2] = color[2]
        seg_single_class[seg_single_class < cl] = 0
        seg_single_class[seg_single_class > cl] = 0
        seg_single_class[seg_single_class == cl] = 255
        sss = cv2.add(tmp_image, np.zeros(np.shape(tmp_image),
                                          dtype=np.uint8), mask=seg_single_class)
        ret_img = cv2.addWeighted(ret_img, 0.8, sss, 0.2, 0)
    return ret_img


def max_region_filter(image, area_th):
    _, binary = cv2.threshold(image, 0.9, 255, cv2.THRESH_BINARY)
    contours, hierarch = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    out = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
    max_area = 0
    max_index = 0
    if len(contours) == 0:
        print("no contours found")
        return False, out

    for i in range(len(contours)):
        area = cv2.contourArea(contours[i])
        # print("area:", i, area)
        if area > max_area:
            max_area = area
            max_index = i

    out = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
    if max_area > area_th:
        cv2.fillPoly(out, [contours[max_index]], 255)
        # cv2.drawContours(out, contours, max_index, 255)
    else:
        print("contours area is small")
        return False, out
    return True, out


def gen_masked_image(mask, color_image_org, depth_image_org):
    color_image = color_image_org.copy()
    depth_image = depth_image_org.copy()
    color_image = cv2.add(color_image,
                          np.zeros(np.shape(color_image),
                                   dtype=np.uint8), mask=mask)
    depth_image = depth_image[:, :, np.newaxis]
    depth_image = cv2.add(depth_image,
                          np.zeros(np.shape(depth_image), dtype=np.uint16), mask=mask)
    depth_filter(depth_image)
    return color_image, depth_image


def depth_filter(depth_image):
    # 16uint 统计10cm-60cm 的有效深度
    # so multiply 10 + 100
    depth_image[depth_image > 1000] = 0
    hist = cv2.calcHist([depth_image], [0], None, [50], [100, 600])
    # 从小到大排序
    x = np.argsort(hist, axis=0, kind='quicksort', order=None)
    d = x[-1][0]
    # 找到深度比集中分布要深，的第一个点少的区域作为截止深度
    # for i in range(len(x)):
    #     if hist[x[-i]] < 1000 and x[-i] > x[-1]:
    #         d = x[-i]
    #         break
    depth = d * 10 + 100
    if np.count_nonzero(depth_image) ==0:
        mean = 0
    else:
        mean = np.sum(depth_image) / np.count_nonzero(depth_image)
    depth_image[depth_image < depth - 30] = 0
    depth_image[depth_image > depth + 30] = mean
    # depth_image[depth_image > depth + 100] = 0

    kernel = np.ones([5, 5], np.float32) / 25  # 除以25是防止数值溢出
    # cv2.filter2D(depth_image, -1, kernel, dst=depth_image)
    return depth

